import data_process
import torch
from torch.utils.data import DataLoader
import os

path = os.path.dirname(__file__)


class data_loader(torch.utils.data.Dataset):
    def __init__(self, data, label):
        self.data = torch.from_numpy(data)
        self.label = torch.from_numpy(label)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        genotype = self.data[index].float()
        label = self.label[index].float()
        return genotype, label


def outpath(work):
    # create runs/work/ path at first, if already exists, ignore this
    os.makedirs(f'{path}/runs/{work}/', exist_ok=True)
    # ergodic names of all folders
    folder_names = os.listdir(f'{path}/runs/{work}/')
    # look for the biggest index number
    if folder_names:
        number_list = [int(folder_name[len(work):]) for folder_name in folder_names if folder_name.startswith(work)]
        number = max(number_list) + 1
    else:
        number = 1
    new_path = f'{path}/runs/{work}/{work}{number}'
    os.makedirs(new_path)
    print(f"{work}ing results will be saved in {new_path}")
    return new_path


d = data_process.DataProcess('data/train_example.vcf', 'data/train_example.csv')
p_trait_dic, n_trait_dic = d.convert_trait('test/')

p_trait_list, n_trait_list = list(p_trait_dic.keys()), list(n_trait_dic.keys())
device = 'cuda' if torch.cuda.is_available() else 'cpu'

print(n_trait_list)
for trait in p_trait_list:
    print(trait)
    train_data, train_label, test_data, test_label = d.to_dataset(trait, percentage=0.7, is_quality=True)
    train_dataloader = DataLoader(data_loader(train_data, train_label), batch_size=20, shuffle=True,
                                  num_workers=8)

    print(train_data[0,0,0])
    print(train_data[0, 0, 1])
    print(train_data[0,0])
    print(len(train_data[0, 0, 0]))
    print(len(train_data[0, 0, 1]))
    for num_data, (genomap, target_trait) in enumerate(train_dataloader):
        genomap, target_trait = genomap.to(device), target_trait.to(device)
        print(len(genomap))
